"""
AI测试用例生成智能体 - Streamlit前端应用
基于LangChain、百度OCR、硅基流动大模型
"""

import streamlit as st
import os
from pathlib import Path
from datetime import datetime
import logging

# 导入自定义模块
from config import get_config
from utils.logger import setup_logger
from utils.file_utils import (
    save_uploaded_file,
    get_file_extension,
    validate_file_extension,
    get_file_size_mb,
    export_to_excel,
    export_to_xmind,
    clean_html_tags
)
from utils.ocr_usage_tracker import get_today_usage
from utils.baidu_ocr import create_baidu_ocr_tool
from utils.siliconflow_llm import create_siliconflow_llm
from chains.test_case_chain import create_test_case_chain, save_test_cases

# 设置页面配置
st.set_page_config(
    page_title="AI测试用例生成智能体",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 初始化日志
logger = setup_logger(
    name="streamlit_app",
    level=get_config("app.log_level", "INFO")
)

logger.info("=" * 50)
logger.info("应用启动")
logger.info("=" * 50)


def init_session_state():
    """初始化会话状态"""
    if 'generated_test_cases' not in st.session_state:
        st.session_state.generated_test_cases = ""
    if 'requirement_analysis' not in st.session_state:
        st.session_state.requirement_analysis = ""
    if 'processing' not in st.session_state:
        st.session_state.processing = False
    if 'uploaded_file_path' not in st.session_state:
        st.session_state.uploaded_file_path = None
    if 'extracted_text' not in st.session_state:
        st.session_state.extracted_text = ""


def validate_config() -> bool:
    """
    验证配置是否完整
    
    Returns:
        配置是否有效
    """
    # 检查百度OCR配置
    baidu_api_key = get_config("baidu_ocr.api_key")
    if not baidu_api_key or baidu_api_key == "your_baidu_api_key":
        st.error("⚠️ 请先在 config/config.yaml 中配置百度OCR的API密钥")
        return False
    
    # 检查硅基流动配置
    silicon_api_key = get_config("siliconflow.api_key")
    if not silicon_api_key or silicon_api_key == "your_siliconflow_api_key":
        st.error("⚠️ 请先在 config/config.yaml 中配置硅基流动的API密钥")
        return False
    
    return True


def process_uploaded_file(uploaded_file) -> tuple:
    """
    处理上传的文件
    
    Args:
        uploaded_file: Streamlit上传的文件对象
        
    Returns:
        (是否成功, 文件路径或错误信息)
    """
    try:
        # 验证文件扩展名
        allowed_extensions = get_config("upload.allowed_extensions", [])
        if not validate_file_extension(uploaded_file.name, allowed_extensions):
            return False, f"不支持的文件格式。允许的格式: {', '.join(allowed_extensions)}"
        
        # 验证文件大小
        file_bytes = uploaded_file.read()
        file_size_mb = get_file_size_mb(file_bytes)
        max_size_mb = get_config("upload.max_file_size", 10)
        
        if file_size_mb > max_size_mb:
            return False, f"文件过大 ({file_size_mb:.2f}MB)，最大允许 {max_size_mb}MB"
        
        # 保存文件
        upload_dir = get_config("upload.upload_folder", "uploads")
        file_path = save_uploaded_file(
            file_bytes,
            uploaded_file.name,
            upload_dir
        )
        
        logger.info(f"文件上传成功: {file_path}")
        return True, file_path
        
    except Exception as e:
        logger.error(f"文件处理失败: {e}")
        return False, f"文件处理失败: {str(e)}"


def extract_text_from_file(file_path: str) -> tuple:
    """
    从文件中提取文本
    
    Args:
        file_path: 文件路径
        
    Returns:
        (是否成功, 提取的文本或错误信息)
    """
    try:
        file_ext = get_file_extension(file_path)
        
        # 图片和PDF使用OCR
        if file_ext in ['png', 'jpg', 'jpeg', 'pdf']:
            logger.info("使用OCR识别文件...")
            
            # 创建OCR工具
            ocr_tool = create_baidu_ocr_tool(
                app_id=get_config("baidu_ocr.app_id"),
                api_key=get_config("baidu_ocr.api_key"),
                secret_key=get_config("baidu_ocr.secret_key"),
                ocr_type=get_config("baidu_ocr.ocr_type", "accurate_basic")
            )
            
            # 执行OCR
            text = ocr_tool._run(file_path)
            
            if text.startswith("错误") or text.startswith("OCR识别失败"):
                return False, text
            
            logger.info(f"OCR识别成功，提取文本长度: {len(text)} 字符")
            return True, text
        
        # 纯文本文件直接读取
        elif file_ext in ['txt', 'md']:
            with open(file_path, 'r', encoding='utf-8') as f:
                text = f.read()
            logger.info(f"文本文件读取成功，长度: {len(text)} 字符")
            return True, text
        
        # DOCX文件
        elif file_ext == 'docx':
            try:
                import docx
                doc = docx.Document(file_path)
                text = '\n'.join([para.text for para in doc.paragraphs])
                logger.info(f"DOCX文件读取成功，长度: {len(text)} 字符")
                return True, text
            except ImportError:
                return False, "需要安装python-docx库来处理DOCX文件: pip install python-docx"
        
        else:
            return False, f"不支持的文件类型: {file_ext}"
            
    except Exception as e:
        logger.error(f"文本提取失败: {e}")
        return False, f"文本提取失败: {str(e)}"


def generate_test_cases_from_text(
    requirement_text: str,
    test_types: list,
    include_analysis: bool
) -> dict:
    """
    从需求文本生成测试用例
    
    Args:
        requirement_text: 需求文本
        test_types: 测试类型列表
        include_analysis: 是否包含需求分析
        
    Returns:
        生成结果字典
    """
    try:
        logger.info("开始生成测试用例...")
        
        # 创建大模型实例
        llm = create_siliconflow_llm(
            api_key=get_config("siliconflow.api_key"),
            base_url=get_config("siliconflow.base_url"),
            model=get_config("siliconflow.model"),
            temperature=get_config("siliconflow.temperature", 0.7),
            max_tokens=get_config("siliconflow.max_tokens", 4000)
        )
        
        # 创建测试用例生成链
        test_case_chain = create_test_case_chain(
            llm=llm,
            enable_refinement=False  # MVP版本暂不启用优化
        )
        
        # 执行生成
        result = test_case_chain.run(
            requirement_text=requirement_text,
            test_types=test_types if test_types else None,
            include_analysis=include_analysis
        )
        
        logger.info("测试用例生成完成")
        return result
        
    except Exception as e:
        logger.error(f"测试用例生成异常: {e}")
        return {
            "status": "error",
            "error": str(e),
            "test_cases": "",
            "analysis": ""
        }


def main():
    """主函数"""
    
    # 初始化会话状态
    init_session_state()
    
    # 标题和描述
    st.title("🤖 " + get_config("app.title", "AI测试用例生成智能体"))
    st.markdown(get_config("app.description", "基于LangChain的智能测试用例生成工具"))
    
    # 验证配置
    if not validate_config():
        st.stop()
    
    # 侧边栏配置
    with st.sidebar:
        st.header("⚙️ 配置选项")
        
        # 测试类型选择
        available_test_types = get_config("test_case.test_types", [
            "功能测试", "边界测试", "异常测试", "性能测试", "安全测试"
        ])
        
        selected_test_types = st.multiselect(
            "选择测试类型",
            options=available_test_types,
            default=["功能测试", "边界测试", "异常测试"],
            help="选择需要生成的测试用例类型"
        )
        
        # 是否包含需求分析
        include_analysis = st.checkbox(
            "包含需求分析",
            value=False,
            help="是否在生成测试用例前先进行需求分析"
        )
        
        st.divider()
        
        # 模型配置显示
        with st.expander("🔧 当前配置"):
            st.text(f"大模型: {get_config('siliconflow.model')}")
            st.text(f"OCR类型: {get_config('baidu_ocr.ocr_type')}")
            st.text(f"温度: {get_config('siliconflow.temperature')}")
        
        st.divider()
        
        # OCR使用情况（不太显眼的位置）
        ocr_usage = get_today_usage()
        usage_percent = (ocr_usage['count'] / ocr_usage['limit']) * 100
        
        # 根据使用情况选择颜色
        if usage_percent >= 90:
            color = "#ff4444"  # 红色
        elif usage_percent >= 70:
            color = "#ff9800"  # 橙色
        else:
            color = "#888888"  # 灰色
        
        st.caption(
            f"<span style='color: {color}; font-size: 11px;'>OCR今日: {ocr_usage['count']}/{ocr_usage['limit']}</span>",
            unsafe_allow_html=True
        )
        
        # 版本信息
        st.caption(f"版本: {get_config('app.version', '1.0.0')}")
    
    # 主界面
    st.header("📝 输入需求文档")
    
    # 创建两个标签页：文件上传和文本输入
    tab1, tab2 = st.tabs(["📁 文件上传", "✏️ 文本输入"])
    
    requirement_text = ""
    
    with tab1:
        st.markdown("**支持的格式**: PNG, JPG, JPEG, PDF, TXT, DOCX, MD")
        
        uploaded_file = st.file_uploader(
            "选择需求文档",
            type=get_config("upload.allowed_extensions", []),
            help="上传需求文档、产品文档或功能说明"
        )
        
        if uploaded_file is not None:
            # 显示文件信息
            col1, col2, col3 = st.columns(3)
            col1.metric("文件名", uploaded_file.name)
            col2.metric("文件大小", f"{get_file_size_mb(uploaded_file.getvalue()):.2f} MB")
            col3.metric("文件类型", get_file_extension(uploaded_file.name).upper())
            
            # 处理文件
            if st.button("🔍 提取文本", type="primary", key="extract_btn"):
                # 检查OCR配额
                file_ext = get_file_extension(uploaded_file.name)
                if file_ext in ['png', 'jpg', 'jpeg', 'pdf']:
                    ocr_usage = get_today_usage()
                    if ocr_usage['count'] >= ocr_usage['limit']:
                        st.error(f"❌ OCR配额已用完（{ocr_usage['count']}/{ocr_usage['limit']}），请明天再试或使用文本文件")
                        st.stop()
                    elif ocr_usage['remaining'] <= 10:
                        st.warning(f"⚠️ OCR配额即将用完，今日剩余 {ocr_usage['remaining']} 次")
                
                with st.spinner("正在处理文件..."):
                    # 保存文件
                    success, result = process_uploaded_file(uploaded_file)
                    
                    if success:
                        st.session_state.uploaded_file_path = result
                        st.success(f"✅ 文件保存成功: {Path(result).name}")
                        
                        # 提取文本
                        success, text = extract_text_from_file(result)
                        
                        if success:
                            # 保存提取的文本到session_state
                            st.session_state.extracted_text = text
                            st.success(f"✅ 文本提取成功，共 {len(text)} 字符")
                            st.rerun()
                        else:
                            st.error(f"❌ {text}")
                    else:
                        st.error(f"❌ {result}")
            
            # 如果有提取的文本，显示文本和生成按钮
            if st.session_state.extracted_text:
                st.divider()
                st.subheader("📄 提取的文本内容")
                
                # 显示提取的文本
                with st.expander("查看完整文本", expanded=False):
                    st.text_area(
                        "提取内容",
                        value=st.session_state.extracted_text,
                        height=300,
                        key="show_extracted_text",
                        disabled=True
                    )
                
                st.info(f"📊 文本长度: {len(st.session_state.extracted_text)} 字符")
                
                # 添加直接生成测试用例的按钮
                col1, col2 = st.columns([3, 1])
                with col1:
                    if st.button("🚀 根据提取的文本生成测试用例", type="primary", use_container_width=True, key="generate_from_ocr"):
                        if not selected_test_types:
                            st.warning("⚠️ 请在侧边栏至少选择一种测试类型")
                        else:
                            st.session_state.processing = True
                            
                            with st.spinner("🔄 正在生成测试用例，请稍候..."):
                                result = generate_test_cases_from_text(
                                    requirement_text=st.session_state.extracted_text,
                                    test_types=selected_test_types,
                                    include_analysis=include_analysis
                                )
                            
                            st.session_state.processing = False
                            
                            if result["status"] == "success":
                                # 清理HTML标签后再保存
                                test_cases_text = result["test_cases"]
                                test_cases_text = clean_html_tags(test_cases_text)
                                
                                st.session_state.generated_test_cases = test_cases_text
                                if include_analysis:
                                    st.session_state.requirement_analysis = result["analysis"]
                                
                                st.success("✅ 测试用例生成成功！")
                                st.rerun()
                            else:
                                st.error(f"❌ 生成失败: {result.get('error', '未知错误')}")
                
                with col2:
                    if st.button("🗑️ 清空", use_container_width=True, key="clear_extracted"):
                        st.session_state.extracted_text = ""
                        st.rerun()
    
    with tab2:
        requirement_text = st.text_area(
            "直接输入或粘贴需求文档内容",
            height=300,
            placeholder="请输入需求描述、功能说明或产品需求文档...",
            help="支持粘贴文本内容"
        )
    
    # 生成测试用例按钮
    st.divider()
    
    if requirement_text:
        st.info(f"📊 当前需求文本长度: {len(requirement_text)} 字符")
        
        col1, col2, col3 = st.columns([2, 1, 1])
        
        with col1:
            generate_btn = st.button(
                "🚀 生成测试用例",
                type="primary",
                use_container_width=True,
                disabled=st.session_state.processing
            )
        
        with col2:
            if st.session_state.generated_test_cases:
                clear_btn = st.button("🗑️ 清空结果", use_container_width=True)
                if clear_btn:
                    st.session_state.generated_test_cases = ""
                    st.session_state.requirement_analysis = ""
                    st.rerun()
        
        # 生成测试用例
        if generate_btn:
            if not selected_test_types:
                st.warning("⚠️ 请至少选择一种测试类型")
            else:
                st.session_state.processing = True
                
                with st.spinner("🔄 正在生成测试用例，请稍候..."):
                    result = generate_test_cases_from_text(
                        requirement_text=requirement_text,
                        test_types=selected_test_types,
                        include_analysis=include_analysis
                    )
                
                st.session_state.processing = False
                
                if result["status"] == "success":
                    # 清理HTML标签后再保存
                    test_cases_text = result["test_cases"]
                    # 清理表格内容中的HTML标签
                    test_cases_text = clean_html_tags(test_cases_text)
                    
                    st.session_state.generated_test_cases = test_cases_text
                    if include_analysis:
                        st.session_state.requirement_analysis = result["analysis"]
                    
                    st.success("✅ 测试用例生成成功！")
                    st.rerun()
                else:
                    st.error(f"❌ 生成失败: {result.get('error', '未知错误')}")
    
    # 显示生成结果
    if st.session_state.generated_test_cases:
        st.header("📋 生成的测试用例")
        
        # 如果有需求分析，先显示
        if st.session_state.requirement_analysis:
            with st.expander("🔍 需求分析", expanded=False):
                st.markdown(st.session_state.requirement_analysis)
        
        # 显示测试用例
        st.markdown(st.session_state.generated_test_cases)
        
        # 导出功能
        st.divider()
        st.subheader("📤 导出测试用例")
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 第一行：Markdown 和本地保存
        col1, col2 = st.columns(2)
        
        with col1:
            # 导出为Markdown
            filename_md = f"test_cases_{timestamp}.md"
            
            st.download_button(
                label="📥 下载 Markdown 格式",
                data=st.session_state.generated_test_cases,
                file_name=filename_md,
                mime="text/markdown",
                use_container_width=True
            )
        
        with col2:
            # 保存到本地
            if st.button("💾 保存到本地 (MD)", use_container_width=True):
                output_dir = Path("outputs")
                output_path = output_dir / filename_md
                
                if save_test_cases(st.session_state.generated_test_cases, str(output_path)):
                    st.success(f"✅ 已保存到: {output_path}")
                else:
                    st.error("❌ 保存失败")
        
        # 第二行：Excel 和 XMind
        col3, col4 = st.columns(2)
        
        with col3:
            # 导出为Excel
            try:
                excel_file = export_to_excel(st.session_state.generated_test_cases)
                if excel_file:
                    filename_excel = f"test_cases_{timestamp}.xlsx"
                    st.download_button(
                        label="📊 下载 Excel 格式",
                        data=excel_file,
                        file_name=filename_excel,
                        mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                        use_container_width=True
                    )
                else:
                    if st.button("📊 下载 Excel 格式", use_container_width=True, disabled=True):
                        st.error("❌ Excel导出功能不可用")
            except Exception as e:
                if st.button("📊 下载 Excel 格式", use_container_width=True, disabled=True):
                    st.error(f"❌ Excel导出失败: {str(e)}")
        
        with col4:
            # 导出为XMind
            try:
                xmind_file = export_to_xmind(st.session_state.generated_test_cases)
                if xmind_file:
                    filename_xmind = f"test_cases_{timestamp}.xmind"
                    st.download_button(
                        label="🧠 下载 XMind 格式",
                        data=xmind_file,
                        file_name=filename_xmind,
                        mime="application/xmind",
                        use_container_width=True
                    )
                else:
                    st.button("🧠 下载 XMind 格式", use_container_width=True, disabled=True)
                    st.error("❌ XMind导出失败：请检查日志查看详细信息")
            except Exception as e:
                st.button("🧠 下载 XMind 格式", use_container_width=True, disabled=True)
                st.error(f"❌ XMind导出失败: {str(e)}")
    
    # 页脚
    st.divider()
    st.markdown("""
    <div style='text-align: center; color: gray; font-size: 12px;'>
        <p>💡 提示：生成的测试用例仅供参考，请根据实际情况调整</p>
        <p>🔧 技术栈：LangChain + 百度OCR + 硅基流动大模型 + Streamlit</p>
    </div>
    """, unsafe_allow_html=True)


if __name__ == "__main__":
    main()

